package com.linkhub.oms.component.trace.core;

import cn.hutool.core.util.IdUtil;
import com.alibaba.ttl.TransmittableThreadLocal;
import com.linkhub.oms.common.util.QtBeanUtil;
import com.linkhub.oms.common.util.QtOptionalUtil;
import com.linkhub.oms.component.trace.util.MDCUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.util.Objects;
import java.util.function.Consumer;

import static com.linkhub.oms.common.constant.TraceConstant.REQUEST_ID_CONCAT_SPAN_ID;

/**
 * @Description: TraceContext
 * @Author: echo
 * @Date: 2022/8/31 17:30
 * @Version: 1.0
 */
@Slf4j
public class QtTraceContext {

    /**
     * 分布式追踪信息存储, 基于 {@link com.alibaba.ttl.TransmittableThreadLocal} 实现
     */
    private static final TransferInfoHolder TRACE_INFO = new TransferInfoHolder();

    /**
     * 异步嵌套调用场景的 {@code parentSpanId} 存储
     */
    private static final ThreadLocal<String> PARENT_SPAN_ID = new TransmittableThreadLocal<>();

    /**
     * 基于外部获取到 {@code traceId} 初始化 {@code TraceContext}
     *
     * @param traceId traceId
     */
    public static void initTraceContext(String traceId) {
        if (StringUtils.isBlank(traceId) || !traceId.contains(REQUEST_ID_CONCAT_SPAN_ID)) {
            return;
        }
        String[] requestIdAndParentSpanId = traceId.split(REQUEST_ID_CONCAT_SPAN_ID);
        QtTransferInfo qtTransferInfo = create();
        // 设置新的值
        qtTransferInfo.setRequestId(requestIdAndParentSpanId[0]);
        qtTransferInfo.setParentSpanId(requestIdAndParentSpanId[1]);
        set(qtTransferInfo);
    }

    /**
     * 如果 {@code context} 为空，则基于外部提供的 {@code requestId} 初始化
     *
     * @param requestId requestId
     */
    public static void initIfAbsent(String requestId) {
        QtTransferInfo qtTransferInfo = TRACE_INFO.get();
        if (StringUtils.isBlank(requestId) || Objects.nonNull(qtTransferInfo)) {
            return;
        }
        qtTransferInfo = create();
        qtTransferInfo.setRequestId(requestId);
        set(qtTransferInfo);
    }

    /**
     * 获取一次请求的唯一标识
     *
     * @return requestId
     */
    public static String getRequestId() {
        return getTransferInfo().getRequestId();
    }

    /**
     * 获取 {@code traceId}, 不会计算 {@code spanId}
     *
     * @return traceId
     */
    public static String getTraceId() {
        return getAndConsumerTraceId(null, false);
    }


    /**
     * 获取和消费 {@code traceId}
     *
     * @param traceIdConsumer traceIdConsumer
     * @param incrementSpanId incrementSpanId
     * @return traceId
     */
    public static String getAndConsumerTraceId(Consumer<String> traceIdConsumer, boolean incrementSpanId) {

        String traceId = "";
        try {
            QtTransferInfo qtTransferInfo = getTransferInfo();

            // 必须复制一个对象，否则会影响其他线程获取QtTransferInfo
            QtTransferInfo newTransferInfo = QtBeanUtil.copyProperties(qtTransferInfo, new QtTransferInfo());
            // 兼容异步嵌套调用场景
            QtOptionalUtil.of(getParentSpanId())
                    .ifPresent(newTransferInfo::setParentSpanId);

            // 基于新的对象递增spanId
            if (incrementSpanId) {
                int spanId = newTransferInfo.getAndIncrementSpanId();
                traceId = newTransferInfo.createTraceId(spanId);
            } else {
                // 基于ttl存储的数据创建traceId
                traceId = newTransferInfo.createTraceId();
            }

            if (traceIdConsumer != null) {
                traceIdConsumer.accept(traceId);
            }

            set(qtTransferInfo);
        } catch (Exception e) {
            log.error("getAndConsumerTraceId error", e);
        }

        return traceId;
    }

    /**
     * 新建一个traceId
     *
     * @return traceId
     */
    public static String createTraceId() {
        QtTransferInfo qtTransferInfo = create();
        set(qtTransferInfo);
        return qtTransferInfo.createTraceId();
    }

    public static boolean isEmpty() {
        return Objects.isNull(TRACE_INFO.get());
    }
    /**
     * 清除ttl数据
     */
    public static void remove() {
        remove(true);
    }

    public static void remove(boolean removeMDCTraceId) {
        TRACE_INFO.remove();
        PARENT_SPAN_ID.remove();
        if (removeMDCTraceId) {
            MDCUtil.removeAllTraceId();
        }
    }

    private static void set(QtTransferInfo qtTransferInfo) {
        TRACE_INFO.set(qtTransferInfo);

    }

    private static QtTransferInfo getTransferInfo() {
        return QtOptionalUtil.of(TRACE_INFO.get())
                .orElseGet(() -> {
                    QtTransferInfo qtTransferInfo = create();
                    set(qtTransferInfo);
                    return qtTransferInfo;
                });
    }

    private static QtTransferInfo create() {
        QtTransferInfo qtTransferInfo = new QtTransferInfo();
        String requestId = IdUtil.fastSimpleUUID().substring(0, 10);
        qtTransferInfo.setRequestId(requestId);
        qtTransferInfo.setLeafSpanIdInfo(new LeafSpanIdInfo());
        return qtTransferInfo;
    }

    private static String getParentSpanId() {
        return PARENT_SPAN_ID.get();
    }

    private static void setParentSpanId(String parentSpanId) {
        PARENT_SPAN_ID.set(parentSpanId);
    }


    private static class TransferInfoHolder extends TransmittableThreadLocal<QtTransferInfo> {

        @Override
        protected void beforeExecute() {
            try {
                // 当异步方法调用异步方法时，获取并设置最新的parentSpanId
                String traceId = MDCUtil.getAndSetSpanTraceId();
                // 设置新key,给该线程的子线程获取spanId
                setParentSpanId(traceId.split(REQUEST_ID_CONCAT_SPAN_ID)[1]);
            } catch (Exception e) {
                log.error("TransferInfoHolder beforeExecute error", e);
            }

        }

    }


}
